//      build_new_fom.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp> 
#include <boost/foreach.hpp>

#include "helper/stl_helper.h"
#include "helper/ublas_helper.h"
#include "helper/ml_data.h"
#include "helper/rand_helper.h"
#include "helper/dna_processing.h"
#include "helper/suffix_array.h"

#include "cds_vom.h"

extern boost::program_options::variables_map vm;

boost::posix_time::ptime now() ;
boost::posix_time::time_duration duration() ;

void compute_likelihood(Sequences& S, const std::string& llout, vom::vom& model)
{
  FileAccessWrite(out, llout);
  int j = 1;
  BOOST_FOREACH(const std::string& s, S) {
    double ll = model.log_likelihood(s);
    out << ">" << j << "\n" << ll << ENDL;
    j++;
  }
}

void compute_likelihood_lambda(Sequences& S, const std::string& llout, vom::vom& model)
{
  FileAccessWrite(out, llout);
  int j = 1;
  BOOST_FOREACH(const std::string& s, S) {
    double ll = model.log_likelihood_lambda(s);
    out << ">" << j << "\n" << ll << ENDL;
    j++;
  }
}

void build_new_fom()
{
  std::string inFile = VM(vm, "input", std::string);
  vi_t orders = VM(vm, "order", std::vector<int>);
  int period = VM(vm, "period", int);
  MSG("++ Build FOM input=" << inFile << " period=" << period);

  // read sequences and sort suffixes
  Sequences S;
  S.fasta_read(inFile);
  MSG("read " << sz(S) << " sequences.");

  // likelihood files
  std::vector<Sequences> llSeqs;
  std::vector<std::string> ll, llout;
  if (vm.count("ll") && vm.count("llout")) {
    ll = VM(vm, "ll", std::vector<std::string>);
    llout = VM(vm, "llout", std::vector<std::string>);
    if (sz(ll) != sz(llout)) 
      FATAL_ERROR("Check --ll and --llout, should have same number of file names");
    for (int i = 0; i < sz(ll); i++) {
      Sequences seqs;
      seqs.fasta_read(ll.at(i));
      llSeqs.pb(seqs);
      MSG("Read likelihood sequences ll=" << ll.at(i) << " llout=" << llout.at(i));
    }
  }

  MSG("All data read duration=" << duration() << " now=" << now());
  BOOST_FOREACH(int order, orders) {
    vom::vom model("ACGT", period);
    model.build_fom_flat(S, order);
    MSG("built FOM order=" << order << " duration=" << duration() << " now=" << now());

    if (vm.count("fom_stats")) {
      std::string statsFile = VM(vm, "fom_stats", std::string) + "." + tos(order);
      model.compute_stats();
      model.print_stats(statsFile);
    }

    if (vm.count("fom_ll")  && sz(llSeqs) > 0) { 
      for (int i = 0; i < sz(llSeqs); i++) {
	str_t outFile = llout.at(i) + "." + tos(order);
	compute_likelihood(llSeqs[i], outFile, model);
      }
      MSG("FOM likelihood duration=" << duration() << " now=" << now());
    }
    
    if (vm.count("prune_ll") && vm.count("eta")) {
      std::vector<double> etas = VM(vm, "eta", std::vector<double>);
      BOOST_FOREACH(double eta, etas) {
	vom::vom pruned_model("ACGT", period);
	pruned_model.copy(model);
	pruned_model.prune(eta);

	if (vm.count("fom_stats")) {
	  std::string statsFile = VM(vm, "fom_stats", std::string) 
	    + ".prune." + tos(eta) + "." + tos(order);
	  pruned_model.compute_stats();
	  pruned_model.print_stats(statsFile);
	}
	for (int i = 0; i < sz(llSeqs); i++) {
	  str_t outFile = llout.at(i) + ".prune." + tos(eta) + "." + tos(order);
	  compute_likelihood(llSeqs[i], outFile, pruned_model);
	}
	MSG("Pruned likelihood eta="<<eta << " duration="<<duration());
      }
    }

    if (vm.count("imm_chi_ll") && sz(llSeqs) > 0) {
      MSG("IMM chi-square likelihood");

      vi_t Ts = VM(vm, "imm_Ts", vi_t);
      BOOST_FOREACH(int T, Ts) {
	vom::vom imm_chi_model("ACGT", period);
	imm_chi_model.copy(model);
	imm_chi_model.compute_lambda_chi_square(T);

	if (vm.count("fom_stats")) {
	  std::string statsFile = VM(vm, "fom_stats", std::string) 
	    + ".chi." + tos(T) + "." + tos(order);
	  imm_chi_model.compute_stats_lambda();
	  imm_chi_model.print_stats(statsFile);
	}

	for (int i = 0; i < sz(llSeqs); i++) {
	  str_t outFile = llout.at(i) + ".chi." + tos(T) + "." + tos(order);
	  compute_likelihood_lambda(llSeqs[i], outFile, imm_chi_model);
	}
	MSG("IMM chi-square likelihood T="<<T << " duration="<<duration());	
      }
    }

    if (vm.count("imm_di_ll") && sz(llSeqs) > 0) {
      MSG("IMM deleted interpolation likelihood");

      std::vector<double> deltas = VM(vm, "imm_delta", std::vector<double>);
      str_t validationFile = VM(vm, "validation", std::string);
      Sequences VS;
      VS.fasta_read(validationFile);
      MSG("read " << sz(VS) << " validation sequences.");
      BOOST_FOREACH(double delta, deltas) {
	vom::vom imm_di_model("ACGT", period);
	imm_di_model.copy(model);
	imm_di_model.compute_lambda_deleted_interpolation(VS, delta);

	if (vm.count("fom_stats")) {
	  std::string statsFile = VM(vm, "fom_stats", std::string) 
	    + ".di." + tos(delta) + "." + tos(order);
	  imm_di_model.compute_stats_lambda();
	  imm_di_model.print_stats(statsFile);
	}

	for (int i = 0; i < sz(llSeqs); i++) {
	  str_t outFile = llout.at(i) + ".di." + tos(delta) + "." + tos(order);
	  compute_likelihood_lambda(llSeqs[i], outFile, imm_di_model);
	}
	MSG("IMM deleted interpolation likelihood delta="<<delta << " duration="<<duration());	
      }
    }

    if (vm.count("imm_prune_chi_ll") && sz(llSeqs) > 0) {
      MSG("IMM chi-squared then prune likelihood");

      vi_t Ts = VM(vm, "imm_Ts", vi_t);
      BOOST_FOREACH(int T, Ts) {
	vom::vom imm_chi_model("ACGT", period);
	imm_chi_model.copy(model);
	imm_chi_model.compute_lambda_chi_square(T);

	std::vector<double> etas = VM(vm, "eta", std::vector<double>);
	BOOST_FOREACH(double eta, etas) {
	  vom::vom pruned_model("ACGT", period);
	  pruned_model.copy(imm_chi_model);
	  pruned_model.prune_lambda(eta);

	  if (vm.count("fom_stats")) {
	    std::string statsFile = VM(vm, "fom_stats", std::string) 
	      + ".chi.prune." + tos(T) + "." + tos(eta) + "." + tos(order);
	    pruned_model.compute_stats_lambda();
	    pruned_model.print_stats(statsFile);
	  }
	  
	  for (int i = 0; i < sz(llSeqs); i++) {
	    str_t outFile = llout.at(i) + ".chi.prune." + tos(T) + "." + tos(eta) + "." + tos(order);
	    compute_likelihood_lambda(llSeqs[i], outFile, pruned_model);
	  }
	  MSG("IMM chi-square then prune likelihood T="<<T << " eta=" << eta << " duration="<<duration());	
	}
      }
    }
    
    if (vm.count("imm_prune_di_ll") && sz(llSeqs) > 0) {
      MSG("IMM deleted interpolation then prune likelihood");

      std::vector<double> deltas = VM(vm, "imm_delta", std::vector<double>);
      str_t validationFile = VM(vm, "validation", std::string);
      Sequences VS;
      VS.fasta_read(validationFile);
      MSG("read " << sz(VS) << " validation sequences.");
      BOOST_FOREACH(double delta, deltas) {
	vom::vom imm_di_model("ACGT", period);
	imm_di_model.copy(model);
	imm_di_model.compute_lambda_deleted_interpolation(VS, delta);

	std::vector<double> etas = VM(vm, "eta", std::vector<double>);
	BOOST_FOREACH(double eta, etas) {
	  vom::vom pruned_model("ACGT", period);
	  pruned_model.copy(imm_di_model);
	  pruned_model.prune_lambda(eta);

	  if (vm.count("fom_stats")) {
	    std::string statsFile = VM(vm, "fom_stats", std::string) 
	      + ".di.prune." + tos(delta) + "." + tos(eta) + "." + tos(order);
	    pruned_model.compute_stats_lambda();
	    pruned_model.print_stats(statsFile);
	  }
	  
	  for (int i = 0; i < sz(llSeqs); i++) {
	    str_t outFile = llout.at(i) + ".di.prune." + tos(delta) + "." + tos(eta) + "." + tos(order);
	    compute_likelihood_lambda(llSeqs[i], outFile, pruned_model);
	  }
	  MSG("IMM deleted interpolation then prune likelihood delta="<<delta << " eta=" << eta << " duration="<<duration());	
	}
      }
    }
  }
}
